#!/usr/bin/env python

# 从 denovo 预测的外显子中找到 exonerate 对应的片段
from BCBio import GFF
import os
import shutil    
import copy
import argparse
import sys


parser = argparse.ArgumentParser(
    description='''对exonerate的结果（需要用exonerate_gtf2gff.py先转化为gff3格式），检查augustus等denovo中对应的预测基因
    用法:
    change_evm_id.py -i nig_trans6.gene_structures_post_PASA_updates.1210085.gff3 -o pasa_update.gff3
    由大天才于2021年7月7日创建于浙江农业大学''')

parser.add_argument('-i',
                help='必须给定，输入的exonerate gff3文件 ')

parser.add_argument('-d',
                help='必须给定，输入的denovo gff3文件 ')

parser.add_argument('-o',
                help='必须给定，输出文件的路径，会删除原有文件，创建一个新文件，所以千万不要用./')

parser.add_argument('-l',action='store_true',
                help='是否只提取最多外显子的哪个 denovo结果')

args = parser.parse_args()

if not args.i or not args.d or not args.o:
    parser.print_help()
    sys.exit()



infile = args.i

outpath = args.o

defile =args.d

if not args.l:
    longest = False
else:
    longest = True




try:
    shutil.rmtree(outpath)
except:
    pass


try:
    os.mkdir(outpath)
except:
    pass


first_exon_dic = {}




# 读取 denovo的结果  建立字典

scaf_de_dic = {}

one_gene_dic = {
}

with open(defile) as fila:
    
    for i in fila:

        k = i.strip().split('\t')

        if len(k)>=8:

            scaf = k[0]

            if scaf not in scaf_de_dic:

                scaf_de_dic[scaf] = []

            if k[2] == 'gene':

                if one_gene_dic !={}:

                    scaf_de_dic[scaf].append(copy.deepcopy(one_gene_dic))

                one_gene_dic = {'gene':k, 'exon':[]}

            elif k[2] =='exon':
                
                one_gene_dic['exon'].append(k)
                
                
                

import re

ID_pattern = re.compile(r'ID=([^;]+);*')
Name_pattern = re.compile(r'Name=([^;]+);*')
Parent_pattern = re.compile(r'Parent=([^;]+);*')





inhandle = open(infile , 'r')


for rec in GFF.parse(inhandle):

    scaf = rec.id

    for feature in rec.features:
        # 找到第一个 exon

        #  每个独立的基因
        if feature.type =='gene':
            
            contain_index = []
            
            strand = str(feature.strand).replace('-1','-').replace('1','+')
            
            exon = ''
            
            for sub_feature in feature.sub_features:
                
                if sub_feature.type == 'exon':
                    
                    exon = sub_feature
                    
                    break
            if exon!= '':
            
                for index in scaf_de_dic[scaf]:
                    
                    # 找到 对应的exon
                    if ((exon.location.start in range(int(index['gene'][3]) , int(index['gene'][4])) 
                        or exon.location.end in range(int(index['gene'][3]) , int(index['gene'][4]))) and 
                       strand ==index['gene'][6]):
                        
                        for index_exon in index['exon']:
                            if ((exon.location.start in range(int(index_exon[3]) , int(index_exon[4])) 
                                or exon.location.end in range(int(index_exon[3]) , int(index_exon[4]))) and 
                                strand ==index_exon[6]):
                                
                                # 检查对应的 exon 是否 移码 一致
                                if (int(index_exon[3]) - exon.location.start-1)%3 == 0 :
                                    
                                    contain_index.append(index)
                                    
                                    break
      
            writehandle = open(outpath+'/'+feature.id+'.gff','w')

            writehandle.write('#'+feature.id+'\n')

            if contain_index == []:
                
                print("failed to find matched denovo annotation for %s" % feature.id)
            
            else:
                
                

                
                if longest:
                    contain_index = [sorted(contain_index, key = lambda x:-len(x['exon']))[0]]
                    
                #print(contain_index)

            
                for index in contain_index:
        
                    
                    #print(index)
                    
                    writehandle.write('\t'.join(index['gene'])+'\n')
                    
                    ID = ID_pattern.findall(index['gene'][8])[0]
                    Name = Name_pattern.findall(index['gene'][8])[0]

                    
                    mRNA = index['gene']
                    mRNA[2] = 'mRNA'
                    mRNA[8] = 'ID='+ID+'.mRNA;'+'Name='+Name+'.mRNA;'+'Parent='+ID
                    
                    writehandle.write('\t'.join(mRNA)+'\n')
                    
                    # 加一个mRNA
                    n = 1
                    for e in index['exon']:
                        
                        exon_ID = ID +'.exon' +str(n)
                        
                        exon_name = exon_ID
                        
                        parent = ID+'.mRNA;'
                        
                        e[8] = 'ID='+exon_ID+';Name='+exon_name+';Parent='+parent
                        
                        line = '\t'.join(e) 
                        
                        writehandle.write(line+'\n'+line.replace('exon','CDS')+'\n')
                    
                    
                writehandle.close()